import nebulaPool from 'k6/x/nebulagraph';
import { check } from 'k6';
import { Trend } from 'k6/metrics';
import { sleep } from 'k6';

var lantencyTrend = new Trend('latency');
var responseTrend = new Trend('responseTime');
var rowSizeTrend = new Trend('rowSize');
// initial nebula connect pool
var pool = nebulaPool.initWithSize("{{ address }}", 2000, 200*{{ vu }});

// initial session for every vu
var session = pool.getSession("{{ user }}", "{{ password }}");
session.execute("USE {{ space }}");

String.prototype.format = function () {
  var formatted = this;
  var data = arguments[0];

  formatted = formatted.replace(/\{(\d+)\}/g, function(match, key){
    return data[key]
  });
  return formatted;
};

export function setup() {
  // config csv file
  pool.configCSV("{{ csv_path }}", "|", false);
  // config output file, save every query information
  pool.configOutput("{{ output_path }}");
  sleep(1)
}

export default function (data) {
  let ngql = '{{ nGQL }}';
  let batches = [];
  let batchSize = 100;
  // batch size 100
  for (let i = 0; i < batchSize; i++) {
    let d = session.getData();
    // {0} means the first column data in the csv file
    let value =  "{{ value }}".format(d);
    batches.push(value);
  };
  ngql = ngql + " " + batches.join(',');
  let response = session.execute(ngql);
  check(response, {
    "IsSucceed": (r) => r.isSucceed() === true
  });
  // add trend
  lantencyTrend.add(response.getLatency());
  responseTrend.add(response.getResponseTime());
  rowSizeTrend.add(response.getRowSize());
};

export function teardown() {
  pool.close();
}
